#pragma once
#include <cassert>
#include <iostream>

namespace ns_RBTree
{
    enum COLOR
    {
        BLACK,
        RED
    };

    template <class T>
    struct RBTreeNode
    {
    public:
        T _key;
        COLOR _clr;
        RBTreeNode<T> *_left;
        RBTreeNode<T> *_right;
        RBTreeNode<T> *_parent;

        RBTreeNode(const T &key = T()) : _key(key), _clr(RED), _left(nullptr), _right(nullptr), _parent(nullptr)
        {
        }
    };

    template <class T>
    class RBTree
    {
    public:
        typedef RBTreeNode<T> Node;

        RBTree() : _root(nullptr)
        {
        }

        bool insert(const T &key)
        {
            if (_root == nullptr)
            {
                _root = new Node(key);
                _root->_clr = BLACK;
                return true;
            }
            Node *cur = _root;
            Node *parent = nullptr;
            while (cur)
            {
                parent = cur;
                if (cur->_key > key)
                    cur = cur->_left;
                else if (cur->_key < key)
                    cur = cur->_right;
                else
                    return false;
            }
            cur = new Node(key);
            if (parent->_key > key)
                parent->_left = cur;
            else
                parent->_right = cur;
            cur->_parent = parent;

            // 调整平衡
            while (parent && parent->_clr == RED)
            {
                Node *grandParent = parent->_parent;
                if (grandParent->_left == parent)
                {
                    Node *uncle = grandParent->_right;
                    if (uncle && uncle->_clr == RED)
                    {
                        parent->_clr = BLACK;
                        uncle->_clr = BLACK;
                        grandParent->_clr = RED;

                        cur = grandParent;
                        parent = cur->_parent;
                    }
                    else
                    {
                        if(parent->_left == cur)
                        {
                            RotateR(grandParent);
                            grandParent->_clr = RED;
                            parent->_clr = BLACK;
                        }
                        else
                        {
                            RotateL(parent);
                            RotateR(grandParent);
                            grandParent->_clr = RED;
                            cur->_clr = BLACK;
                        }
                        break;
                    }
                }
                else
                {
                    Node *uncle = grandParent->_left;
                    if (uncle && uncle->_clr == RED)
                    {
                        parent->_clr = BLACK;
                        uncle->_clr = BLACK;
                        grandParent->_clr = RED;

                        cur = grandParent;
                        parent = cur->_parent;
                    }
                    else
                    {
                        if(parent->_right == cur)
                        {
                            RotateL(grandParent);
                            grandParent->_clr = RED;
                            parent->_clr = BLACK;
                        }
                        else
                        {
                            RotateR(parent);
                            RotateL(grandParent);
                            grandParent->_clr = RED;
                            cur->_clr = BLACK;
                        }
                        break;
                    }
                }
            }
            _root->_clr = BLACK;

            return true;
        }

        void Inorder()
        {
            _Inorder(_root);
            std::cout << std::endl;
        }

        bool Blance()
        {
            if(_root && _root->_clr == RED)
            {
                std::cout << "根节点是红色的" << std::endl;
                return false;
            }
            int Num = BlackNum();
            return _Blance(_root, Num);
        }

        int BlackNum()
        {
            int cnt = 0;
            Node* cur = _root;
            while(cur)
            {
                if(cur->_clr == BLACK)
                    ++cnt;
                cur = cur->_left;
            }
            return cnt;
        }
    private:
        bool _Blance(Node* root, int Num)
        {
            if(root == nullptr)
                return true;
            if(root->_clr == BLACK)
                --Num;
            if(root->_clr == RED)
            {
                if(root->_parent && root->_parent->_clr == RED)
                {
                    std::cout << "出现了连续的红节点" << std::endl;
                    return false;
                }
            }
            if(root->_left == nullptr && root->_right == nullptr)
            {
                if(Num != 0)
                {
                    std::cout << "每条路径的黑节点数量不相等" << std::endl;
                    return false;
                }
            }
            return _Blance(root->_left, Num) && _Blance(root->_right, Num);
        }

        void _Inorder(Node *&root)
        {
            if (nullptr == root)
                return;
            _Inorder(root->_left);
            std::cout << root->_key << " ";
            _Inorder(root->_right);
        }

        void RotateR(Node *&root)
        {
            Node *Left = root->_left;
            Node *LeftR = Left->_right;

            root->_left = LeftR;
            if (LeftR)
                LeftR->_parent = root;

            Left->_right = root;
            if (root->_parent == nullptr)
            {
                _root = Left;
                _root->_parent = nullptr;
            }
            else
            {
                if (root->_parent->_left == root)
                    root->_parent->_left = Left;
                else
                    root->_parent->_right = Left;
                Left->_parent = root->_parent;
            }
            root->_parent = Left;
        }

        void RotateL(Node *&root)
        {
            Node *Right = root->_right;
            Node *RightL = Right->_left;

            root->_right = RightL;
            if (RightL)
                RightL->_parent = root;

            Right->_left = root;
            if (root->_parent == nullptr)
            {
                _root = Right;
                _root->_parent = nullptr;
            }
            else
            {
                if (root->_parent->_left == root)
                    root->_parent->_left = Right;
                else
                    root->_parent->_right = Right;
                Right->_parent = root->_parent;
            }
            root->_parent = Right;
        }

    private:
        Node *_root;
    };
}